import os

import tool as tl
import numpy as np
import math


def get_params(path):
    loss_path = path + 'loss/'
    distance_path = path + 'result/'
    loss_dirs = os.listdir(loss_path)
    losslist = []
    for loss_dir in loss_dirs:
        loss_array = tl.read_csv(loss_path + loss_dir)
        losslist_vec = []
        for loss in loss_array:
            losslist_vec.append(float(loss) - float(loss_array[-1]))
        losslist.append(losslist_vec)

    distance_dirs = os.listdir(distance_path)
    distancelist = []
    for distance_dir in distance_dirs:
        distancelist.append(tl.to_float(tl.read_csv(distance_path + distance_dir))[0])

    return tl.get_mean_arr(np.array(losslist)), tl.get_mean_arr(np.array(distancelist))


def log_r_curves(KD):
    kd_loss_vec, kd_dis_vec = get_params(KD)
    res = []
    for i in range(30):
        res.append(kd_loss_vec[i] / math.log(kd_dis_vec[i]))
    return res

input_path = '../weights_params/IE demo val loss/'
class_names = os.listdir(input_path)
KDclass_names = []
draw_matrix = []
for name in class_names:
    if( not 'cnn kd cnn' in name) and not('baseline' in name):
        continue
    print('computing the information effectiveness between [{}] and [standard cnn]'.format(name))
    draw_matrix.append(log_r_curves(input_path + name + '/'))
    KDclass_names.append(name)
tl.pic_make_matrix(draw_matrix, KDclass_names, 'Information effectiveness')